/*
 * LZMA2OutputStream
 *
 * Authors: Lasse Collin <lasse.collin@tukaani.org>
 *          Igor Pavlov <http://7-zip.org/>
 *
 * This file has been put into the public domain.
 * You can do whatever you want with this file.
 */

package org.tukaani.xz;

import java.io.DataOutputStream;
import java.io.IOException;

import org.tukaani.xz.lz.LZEncoder;
import org.tukaani.xz.lzma.LZMAEncoder;
import org.tukaani.xz.rangecoder.RangeEncoder;

class LZMA2OutputStream extends FinishableOutputStream
{
  static final int COMPRESSED_SIZE_MAX = 64 << 10;

  private FinishableOutputStream out;

  private final DataOutputStream outData;

  private final LZEncoder lz;

  private final RangeEncoder rc;

  private final LZMAEncoder lzma;

  private final int props; // Cannot change props on the fly for now.

  private boolean dictResetNeeded = true;

  private boolean stateResetNeeded = true;

  private boolean propsNeeded = true;

  private int pendingSize = 0;

  private boolean finished = false;

  private IOException exception = null;


  private static int getExtraSizeBefore(int dictSize)
  {
    return COMPRESSED_SIZE_MAX > dictSize ? COMPRESSED_SIZE_MAX - dictSize : 0;
  }


  static int getMemoryUsage(LZMA2Options options)
  {
    // 64 KiB buffer for the range encoder + a little extra + LZMAEncoder
    int dictSize = options.getDictSize();
    int extraSizeBefore = getExtraSizeBefore(dictSize);
    return 70 + LZMAEncoder.getMemoryUsage(options.getMode(), dictSize, extraSizeBefore, options.getMatchFinder());
  }


  LZMA2OutputStream(FinishableOutputStream out, LZMA2Options options)
  {
    if (out == null)
      throw new NullPointerException();

    this.out = out;
    outData = new DataOutputStream(out);
    rc = new RangeEncoder(COMPRESSED_SIZE_MAX);

    int dictSize = options.getDictSize();
    int extraSizeBefore = getExtraSizeBefore(dictSize);
    lzma =
        LZMAEncoder.getInstance(rc, options.getLc(), options.getLp(), options.getPb(), options.getMode(), dictSize,
            extraSizeBefore, options.getNiceLen(), options.getMatchFinder(), options.getDepthLimit());

    lz = lzma.getLZEncoder();

    byte[] presetDict = options.getPresetDict();
    if (presetDict != null && presetDict.length > 0)
    {
      lz.setPresetDict(dictSize, presetDict);
      dictResetNeeded = false;
    }

    props = (options.getPb() * 5 + options.getLp()) * 9 + options.getLc();
  }

  private static final byte[] BUFFER_WRITE = new byte[1];


  public void write(int b) throws IOException
  {
    BUFFER_WRITE[0] = (byte) b;
    write(BUFFER_WRITE, 0, 1);
  }


  public void write(byte[] buf, int off, int len) throws IOException
  {
    if (off < 0 || len < 0 || off + len < 0 || off + len > buf.length)
      throw new IndexOutOfBoundsException();

    if (exception != null)
      throw exception;

    if (finished)
      throw new XZIOException("Stream finished or closed");

    try
    {
      while (len > 0)
      {
        int used = lz.fillWindow(buf, off, len);
        off += used;
        len -= used;
        pendingSize += used;

        if (lzma.encodeForLZMA2())
          writeChunk();
      }
    } catch (IOException e)
    {
      exception = e;
      throw e;
    }
  }


  private void writeChunk() throws IOException
  {
    int compressedSize = rc.finish();
    int uncompressedSize = lzma.getUncompressedSize();

    assert compressedSize > 0 : compressedSize;
    assert uncompressedSize > 0 : uncompressedSize;

    // +2 because the header of a compressed chunk is 2 bytes
    // bigger than the header of an uncompressed chunk.
    if (compressedSize + 2 < uncompressedSize)
    {
      writeLZMA(uncompressedSize, compressedSize);
    } else
    {
      lzma.reset();
      uncompressedSize = lzma.getUncompressedSize();
      assert uncompressedSize > 0 : uncompressedSize;
      writeUncompressed(uncompressedSize);
    }

    pendingSize -= uncompressedSize;
    lzma.resetUncompressedSize();
    rc.reset();
  }


  private void writeLZMA(int uncompressedSize, int compressedSize) throws IOException
  {
    int control;

    if (propsNeeded)
    {
      if (dictResetNeeded)
        control = 0x80 + (3 << 5);
      else
        control = 0x80 + (2 << 5);
    } else
    {
      if (stateResetNeeded)
        control = 0x80 + (1 << 5);
      else
        control = 0x80;
    }

    control |= (uncompressedSize - 1) >>> 16;
    outData.writeByte(control);

    outData.writeShort(uncompressedSize - 1);
    outData.writeShort(compressedSize - 1);

    if (propsNeeded)
      outData.writeByte(props);

    rc.write(out);

    propsNeeded = false;
    stateResetNeeded = false;
    dictResetNeeded = false;
  }


  private void writeUncompressed(int uncompressedSize) throws IOException
  {
    while (uncompressedSize > 0)
    {
      int chunkSize = Math.min(uncompressedSize, COMPRESSED_SIZE_MAX);
      outData.writeByte(dictResetNeeded ? 0x01 : 0x02);
      outData.writeShort(chunkSize - 1);
      lz.copyUncompressed(out, uncompressedSize, chunkSize);
      uncompressedSize -= chunkSize;
      dictResetNeeded = false;
    }

    stateResetNeeded = true;
  }


  private void writeEndMarker() throws IOException
  {
    assert !finished;

    if (exception != null)
      throw exception;

    lz.setFinishing();

    try
    {
      while (pendingSize > 0)
      {
        lzma.encodeForLZMA2();
        writeChunk();
      }

      out.write(0x00);
    } catch (IOException e)
    {
      exception = e;
      throw e;
    }

    finished = true;
  }


  public void flush() throws IOException
  {
    if (exception != null)
      throw exception;

    if (finished)
      throw new XZIOException("Stream finished or closed");

    try
    {
      lz.setFlushing();

      while (pendingSize > 0)
      {
        lzma.encodeForLZMA2();
        writeChunk();
      }

      out.flush();
    } catch (IOException e)
    {
      exception = e;
      throw e;
    }
  }


  public void finish() throws IOException
  {
    if (!finished)
    {
      writeEndMarker();

      try
      {
        out.finish();
      } catch (IOException e)
      {
        exception = e;
        throw e;
      }

      finished = true;
    }
  }


  public void close() throws IOException
  {
    if (out != null)
    {
      if (!finished)
      {
        try
        {
          writeEndMarker();
        } catch (IOException e)
        {
        }
      }

      try
      {
        out.close();
      } catch (IOException e)
      {
        if (exception == null)
          exception = e;
      }

      out = null;
    }

    if (exception != null)
      throw exception;
  }
}
